
#include "VideoEncoder.h"
#include <memory.h>

VideoEncoder::VideoEncoder(const int fps, const int vWidth, const int vHeight, const int quality, 
						   const int bitrate, const int keyframeshift, const int dropframe, const int fastmode)
{
	m_nFps = fps;
	m_nVideoWidth = vWidth;
	m_nVideoHeight = vHeight;
	m_nQuality = quality * 63 / 100;
	m_nBitRate = bitrate;
	m_nKeyFrameShift = keyframeshift;
	m_bDropFrame = dropframe;
	m_bFastMode = fastmode;

	initTheoraInfo();
	initTheoraEncoder();
}

VideoEncoder::~VideoEncoder(void)
{
	uinitTheoraEncoder();
	free(m_lpEncoderInfo);
	m_lpEncoderInfo = 0;
}


void VideoEncoder::initTheoraInfo()
{
	m_lpEncoderInfo = (th_info*)malloc(sizeof(th_info));
	th_info_init(m_lpEncoderInfo);

	// video size
	m_lpEncoderInfo->pic_width = m_nVideoWidth;
	m_lpEncoderInfo->pic_height = m_nVideoHeight;
	m_lpEncoderInfo->pic_x = 0;
	m_lpEncoderInfo->pic_y = 0;

	// encoded size: multiple of 16
	int nWidth = m_nVideoWidth % 16;
	m_lpEncoderInfo->frame_width = nWidth ? m_nVideoWidth - nWidth + 16 : m_nVideoWidth;
	int nHeight = m_nVideoHeight % 16;
	m_lpEncoderInfo->frame_height = nHeight ? m_nVideoHeight - nHeight + 16 : m_nVideoHeight;

	// color space
	m_lpEncoderInfo->colorspace = TH_CS_UNSPECIFIED;
	//m_lpEncoderInfo->colorspace = TH_CS_ITU_REC_470M; // NTSC
	//m_lpEncoderInfo->colorspace = TH_CS_ITU_REC_470BG; // PAL

	// pixel format
	m_lpEncoderInfo->pixel_fmt = TH_PF_444;
	//m_lpEncoderInfo->pixel_fmt = TH_PF_422;

	// quality (0..63)
	//m_lpEncoderInfo->quality = 63; // maximum
	//m_lpEncoderInfo->quality = 32; // middle
	m_lpEncoderInfo->quality = m_nQuality; // custom

	// key frame shift
	m_lpEncoderInfo->keyframe_granule_shift = m_nKeyFrameShift;

	// frame rate
	m_lpEncoderInfo->fps_numerator = m_nFps;
	m_lpEncoderInfo->fps_denominator = 1;

	// aspect
	m_lpEncoderInfo->aspect_denominator = 0;
	m_lpEncoderInfo->aspect_numerator = 0;

	// bit rate
	m_lpEncoderInfo->target_bitrate = m_nBitRate * 1000;
}

void VideoEncoder::initTheoraEncoder()
{
	// allocate Theora encoder
	m_lpEncoderContext = th_encode_alloc(m_lpEncoderInfo);
	if (0 == m_lpEncoderContext)
		exit(1);

	int arg = 9;
	if (th_encode_ctl(m_lpEncoderContext, TH_ENCCTL_SET_KEYFRAME_FREQUENCY_FORCE, &arg, sizeof(arg)) < 0)
		exit(1);

	//arg = 1;
	//if (th_encode_ctl(m_lpEncoderContext, TH_ENCCTL_SET_VP3_COMPATIBLE, &arg, sizeof(arg)) < 0)
	//	exit(1);

	if (m_bFastMode)
	{
		arg = 0;
		if (th_encode_ctl(m_lpEncoderContext, TH_ENCCTL_GET_SPLEVEL_MAX, &arg, sizeof(arg)) < 0 ||
			th_encode_ctl(m_lpEncoderContext, TH_ENCCTL_SET_SPLEVEL, &arg, sizeof(arg)) < 0)
			exit(1);
	}

	arg = TH_RATECTL_CAP_OVERFLOW;
	if (m_bDropFrame)
		arg |= TH_RATECTL_DROP_FRAMES;
	if (th_encode_ctl(m_lpEncoderContext,TH_ENCCTL_SET_RATE_FLAGS, &arg, sizeof(arg)) < 0)
		exit(1);

	arg = 2 * m_lpEncoderInfo->fps_numerator/m_lpEncoderInfo->fps_denominator;
	if (th_encode_ctl(m_lpEncoderContext, TH_ENCCTL_SET_RATE_BUFFER, &arg, sizeof(arg)) < 0)
		exit(1);

	// init packet
	th_comment comment;
	ogg_packet ogg_pkg;
	th_comment_init(&comment);
	memset(&ogg_pkg, 0, sizeof(ogg_pkg));

	// get first header packet
	m_lpHeaderPackets = 0;
	if (th_encode_flushheader(m_lpEncoderContext, &comment, &ogg_pkg) <= 0)
		return;
	m_lpHeaderPackets = createPacket(&ogg_pkg);

	// get the rest
	PacketList *tmp = m_lpHeaderPackets;
	while (th_encode_flushheader(m_lpEncoderContext, &comment, &ogg_pkg) > 0)
	{
		tmp->next = createPacket(&ogg_pkg);
		tmp = tmp->next;
	}
}

void VideoEncoder::uinitTheoraEncoder()
{
	// free Theora encoder
	th_encode_free(m_lpEncoderContext);
	m_lpEncoderContext = 0;
}

PacketList* VideoEncoder::headerPackets()
{
	PacketList *ret = clonePacket(m_lpHeaderPackets);
	PacketList *tmpDst = ret, *tmpSrc = m_lpHeaderPackets;

	while(tmpDst)
	{
		tmpDst->next = clonePacket(tmpSrc->next);
		tmpDst = tmpDst->next;
		tmpSrc = tmpSrc->next;
	}

	return ret;
}

PacketList * VideoEncoder::remainPackets()
{
	ogg_packet ogg_pkg;

	if (th_encode_packetout(m_lpEncoderContext, 1, &ogg_pkg) <= 0)
		return 0;

	return createPacket(&ogg_pkg);
}

int VideoEncoder::videoWidth()
{
	return m_lpEncoderInfo->pic_width;
}

int VideoEncoder::videoHeight()
{
	return m_lpEncoderInfo->pic_height;
}

PacketList* VideoEncoder::encodeRGB(const unsigned char *data, int last)
{
	th_ycbcr_buffer ycbcr;
	if (createYCbCrBufferFromRGB(data, ycbcr))
		return 0;

	if (th_encode_ycbcr_in(m_lpEncoderContext, ycbcr) < 0)
		exit(1);

	freeYCbCrBuffer(ycbcr);

	return createPacketList(last);
}

PacketList* VideoEncoder::createPacketList(int last)
{
	PacketList *ret = 0, *tmp = 0;;
	ogg_packet ogg_pkg;

	if (th_encode_packetout(m_lpEncoderContext, last, &ogg_pkg) <= 0)
		return 0;

	ret = createPacket(&ogg_pkg);
	tmp = ret;

	while (th_encode_packetout(m_lpEncoderContext, last, &ogg_pkg) > 0)
	{
		tmp->next = createPacket(&ogg_pkg);
		tmp = tmp->next;
	}

	return ret;
}

int VideoEncoder::createYCbCrBufferFromRGB(const unsigned char *data, th_ycbcr_buffer &ycbcr)
{
	if (TH_PF_444 != m_lpEncoderInfo->pixel_fmt)
		return 1;

	int size = m_lpEncoderInfo->pic_width * m_lpEncoderInfo->pic_height;
	unsigned char *y = (unsigned char*)malloc(size);
	unsigned char *u = (unsigned char*)malloc(size);
	unsigned char *v = (unsigned char*)malloc(size);

	// convert from RGB to YUV 444
	rgb_to_yuv444(m_lpEncoderInfo->pic_width, m_lpEncoderInfo->pic_height, data, y, u, v);

	// fill out the ycbcr buffer
	ycbcr[0].width = m_lpEncoderInfo->pic_width;
	ycbcr[0].height = m_lpEncoderInfo->pic_height;
	ycbcr[0].stride = m_lpEncoderInfo->pic_width;
	ycbcr[1].width = m_lpEncoderInfo->pic_width;
	ycbcr[1].height = m_lpEncoderInfo->pic_height;
	ycbcr[1].stride = m_lpEncoderInfo->pic_width;
	ycbcr[2].width = ycbcr[1].width;
	ycbcr[2].height = ycbcr[1].height;
	ycbcr[2].stride = ycbcr[1].stride;

	// fill out data
	ycbcr[0].data = y;
	ycbcr[1].data = u;
	ycbcr[2].data = v;

	return 0;
}

void VideoEncoder::freeYCbCrBuffer(th_ycbcr_buffer buffer)
{
	free(buffer[0].data);
	free(buffer[1].data);
	free(buffer[2].data);
	memset(buffer, 0, sizeof(th_ycbcr_buffer));
}

void VideoEncoder::rgb_to_yuv444(int w, int h, const unsigned char *rgb, unsigned char *y, unsigned char *u, unsigned char *v)
{
	int strike = w * 3;
	for (int i=0; i<h; i++)
	{
		for (int j=0; j<w; j++)
		{
			unsigned char r = *(rgb + i*strike + j * 3 + 2);
			unsigned char g = *(rgb + i*strike + j * 3 + 1);
			unsigned char b = *(rgb + i*strike + j * 3 + 0);
			unsigned char yy = (0.257 * r) + (0.504 * g) + (0.098 * b) + 16;
			unsigned char vv = (0.439 * r) - (0.368 * g) - (0.071 * b) + 128;
			unsigned char uu = -(0.148 * r) - (0.291 * g) + (0.439 * b) + 128;
			*(y + i*w + j) = yy;
			*(u + i*w + j) = uu;
			*(v + i*w + j) = vv;
		}
	}
}
